from paddle.io import DataLoader
from .ShangHaiTech import RTFMTrainDataset, RTFMInferDataset
from .FreatureDataset import FeatureTrainDataset, FeatureInferDataset


def generate_dataset(**kwargs):
    dataset = kwargs.pop('name')
    dataloader = kwargs.pop('dataloader')
    dataset = eval(dataset)(**kwargs)
    dataloader = DataLoader(dataset, **dataloader)
    return dataloader


if __name__ == '__main__':
    train_dataloader = generate_dataset(**{
        "name": "RTFMTrainDataset",
        "data_dir": "D:\\code\\video_abnomal\\shanghaitech\\SH_Train_ten_crop_i3d",
        "target_file": "D:\\code\\video_abnomal\\shanghaitech\\train.csv",
        "dataloader": {
            "batch_size": 16,
            "shuffle": True,
            "drop_last": True,
            "num_workers": 2
        }
    })
    for i, (image, label) in enumerate(train_dataloader()):
        print("="*20 + "%03d" % (i) + "=" * 20)
        print(image.shape)
        print(label.shape)

